# coding=utf-8
from __future__ import print_function, absolute_import
from gm.api import *
import logging
import click
from datetime import datetime, timedelta, date
import pandas as pd
from quant.data import juejin

'''
滚IC吃贴水策略
'''

start_date = '2021-01-01'
end_date = str(datetime.now().date())
ROLL_BY_MA = True
DAYS_TO_DELIST = 10


def init(context):
    schedule(roll, date_rule='1d', time_rule='14:50:00')
    if ROLL_BY_MA:
        context.df_backward = get_backward(days=20)
    context.confirm = True


# 获取当前IC持仓
def get_ic_position(context):
    pos = [p for p in context.account().positions()
           if p.symbol.startswith('CFFEX.IC')]
    assert len(pos) <= 1
    if len(pos) == 0:
        return None
    return pos[0]


def on_order_status(context, order):
    if order.status == OrderStatus_Rejected:
        logging.error('{}: status: {}, target: {}, error: {}'.format(
            context.now.date(), order.status, order.symbol, order.ord_rej_reason_detail))
        click.confirm('Error')


# 获取连续合约贴水的移动平均线
def get_backward(days):
    ics = get_instruments(sec_types=SEC_TYPE_FUTURE, df=True)
    ic2xs = ics[ics.symbol.apply(lambda s: s.startswith('CFFEX.IC2'))]
    ic2xbars = history(','.join(ic2xs.symbol.tolist()), frequency='1d',
                       start_time=date.fromisoformat(start_date) - timedelta(days=days), end_time=end_date, df=True)
    ic2xbars['date'] = ic2xbars.bob.dt.date
    ic2xbars.set_index(['symbol', 'date'], inplace=True)
    zz500 = history(symbol='SHSE.000905', frequency='1d',
                    start_time=start_date,  end_time=end_date, df=True)
    dom_close = []
    sub_dom_close = []
    for dt in zz500.bob.dt.date:
        year, month, days_to_delist = juejin.get_dominant(dt)
        dom = 'CFFEX.IC{}{:02d}'.format(year, month)
        sub_dom = 'CFFEX.IC{}{:02d}'.format(
            *juejin.next_year_month(year, month))
        dom_close.append(ic2xbars.loc[(dom, dt)].close)
        sub_dom_close.append(ic2xbars.loc[(sub_dom, dt)].close)
    df = pd.DataFrame({'date': zz500.bob.dt.date, 'dom_close': dom_close,
                       'sub_dom_close': sub_dom_close, 'index_close': zz500.close}).set_index('date')
    df['backwardation_diff'] = df.dom_close - df.sub_dom_close
    df['backward'] = df.backwardation_diff.rolling(days).mean()
    df.backward.fillna(df.backward[days], inplace=True)
    return df


def roll(context):
    dt = context.now.date()
    year, month, days_to_delist = juejin.get_dominant(dt)
    dominant = 'CFFEX.IC{}{:02d}'.format(year, month)
    sub_dominant = 'CFFEX.IC{}{:02d}'.format(
        *juejin.next_year_month(year, month))
    logging.info('****** %s ******: 主力合约%s' % (dt, dominant))

    zz500 = 'SHSE.000905'
    bars = {bar['symbol']: bar
            for bar in current(symbols=[dominant, zz500, sub_dominant])}
    dom_price = bars[dominant]['price']
    index_price = bars[zz500]['price']
    sub_dom_price = bars[sub_dominant]['price']
    days_to_delist = (get_instrumentinfos(symbols=dominant)[
                      0]['delisted_date'].date() - context.now.date()).days

    p = get_ic_position(context)
    if p is None:
        if dom_price >= index_price:
            logging.info('延迟建仓：主力合约升水, zz500: {:.1f}, {}: {:.1f}, 升水点数: {:.1f}, 升水百分比: {:.2%}'.format(
                index_price, dominant, dom_price, dom_price - index_price, dom_price / index_price - 1))
        else:
            if days_to_delist == 0:
                logging.info('延迟建仓：今日为交割日')
            else:
                logging.info('首次建仓: {}, 贴水点数: {:.1f}, 贴水百分比: {:.2%}'.format(
                    dominant, dom_price - index_price, 1 - dom_price / index_price))
                order_target_volume(
                    symbol=dominant, volume=1, position_side=PositionSide_Long, order_type=OrderType_Market)
                if context.confirm:
                    context.confirm = click.confirm('Continue', default=True)
    else:
        if p.symbol == sub_dominant:
            logging.info('已经换月到次月合约')
        else:
            shouldRoll = False
            if days_to_delist == 0:
                logging.info('交割日移仓换月')
                shouldRoll = True
            elif ROLL_BY_MA:
                diff = dom_price - sub_dom_price
                if diff > context.df_backward.loc[dt].backward:
                    logging.info('基差高于20交易日均值，移仓换月')
                    shouldRoll = True
            elif dom_price >= index_price:
                logging.info('主力合约升水，移仓换月')
                shouldRoll = True
            elif days_to_delist <= DAYS_TO_DELIST:
                logging.info('满足提前移仓换月')
                shouldRoll = True
            if shouldRoll:
                logging.info('移仓换月卖出{}, 买入{}, 点数{:.1f}, 百分比{:.2%}'.format(dominant, sub_dominant, sub_dom_price - index_price, 1 - sub_dom_price / index_price))
                order_target_volume(
                    symbol=dominant, volume=0, position_side=PositionSide_Long, order_type=OrderType_Market)
                order_target_volume(
                    symbol=sub_dominant, volume=1, position_side=PositionSide_Long, order_type=OrderType_Market)
                if context.confirm:
                    context.confirm = click.confirm(
                        'Continue', default=True)
            else:
                logging.info('继续持有主力合约%s' % dominant)


def on_backtest_finished(context, indicator):
    logging.info('*' * 50)
    logging.info(indicator)
    zz500 = history(symbol='SHSE.000905', frequency='1d',
                    start_time=start_date, end_time=end_date, df=True)
    logging.info('zz500: {:.2%}'.format(
        zz500.iloc[-1].close / zz500.iloc[0].close - 1))


if __name__ == '__main__':
    '''
        strategy_id策略ID,由系统生成
        filename文件名,请与本文件名保持一致
        mode实时模式:MODE_LIVE回测模式:MODE_BACKTEST
        token绑定计算机的ID,可在系统设置-密钥管理中生成
        backtest_start_time回测开始时间
        backtest_end_time回测结束时间
        backtest_adjust股票复权方式不复权:ADJUST_NONE前复权:ADJUST_PREV后复权:ADJUST_POST
        backtest_initial_cash回测初始资金
        backtest_commission_ratio回测佣金比例
        backtest_slippage_ratio回测滑点比例
    '''
    logging.getLogger().setLevel(logging.INFO)
    run(strategy_id='f7396dab-ca82-11ec-b6c3-2a16a8610fe9',
        filename='main.py',
        mode=MODE_BACKTEST,
        token='567120cd2f9b2ce8ae3fbd7ee76ed06bf899b512',
        backtest_start_time='%s 09:00:00' % start_date,
        backtest_end_time='%s 15:00:00' % end_date,
        backtest_adjust=ADJUST_PREV,
        backtest_initial_cash=1000000,
        backtest_commission_ratio=0.0001,
        backtest_slippage_ratio=0.0001)
